"""
四川麻将AI模型基础使用示例

本示例展示了如何使用四川麻将AI模型进行以下操作：
1. 准备训练和评估数据
2. 训练模型
3. 评估模型性能
4. 使用模型进行推理
5. 导出模型用于部署
"""
import os
import sys
import numpy as np
import torch
from datetime import datetime

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))

# 导入必要的模块
from sichuanmajiang.model.model import MahjongModel, create_mahjong_model
from sichuanmajiang.model.training_manager import train_model, evaluate_model, export_trained_model
from sichuanmajiang.model.inference import MahjongInference, MahjongDecisionMaker
from sichuanmajiang.data.data_loader import create_data_loader, create_balanced_data_loader
from sichuanmajiang.data.preprocessor import FeatureEngineer
from sichuanmajiang.engine.game_state import MahjongGameState


def create_sample_data(num_samples=1000, feature_dim=324, action_dim=6):
    """
    创建示例训练数据
    
    Args:
        num_samples: 样本数量
        feature_dim: 特征维度
        action_dim: 动作空间大小
    
    Returns:
        features: 特征数据
        actions: 动作标签
        rewards: 奖励值
    """
    print(f"创建示例数据: {num_samples} 样本")
    
    # 生成随机特征（实际应用中应使用真实游戏数据）
    features = np.random.rand(num_samples, feature_dim).astype(np.float32)
    
    # 生成随机动作标签
    actions = np.random.randint(0, action_dim, num_samples)
    
    # 生成随机奖励值
    rewards = np.random.randn(num_samples).astype(np.float32)
    
    # 简单归一化奖励
    rewards = np.tanh(rewards)  # 将奖励缩放到[-1, 1]范围
    
    return features, actions, rewards


def prepare_data():
    """
    准备训练和评估数据
    
    Returns:
        train_loader: 训练数据加载器
        val_loader: 验证数据加载器
        test_loader: 测试数据加载器
    """
    print("准备数据...")
    
    # 创建示例数据
    train_features, train_actions, train_rewards = create_sample_data(num_samples=1000)
    val_features, val_actions, val_rewards = create_sample_data(num_samples=200)
    test_features, test_actions, test_rewards = create_sample_data(num_samples=200)
    
    # 创建数据加载器
    train_loader = create_balanced_data_loader(
        features=train_features,
        actions=train_actions,
        rewards=train_rewards,
        batch_size=32,
        shuffle=True
    )
    
    val_loader = create_data_loader(
        features=val_features,
        actions=val_actions,
        rewards=val_rewards,
        batch_size=32,
        shuffle=False
    )
    
    test_loader = create_data_loader(
        features=test_features,
        actions=test_actions,
        rewards=test_rewards,
        batch_size=32,
        shuffle=False
    )
    
    print(f"数据准备完成:")
    print(f"- 训练数据: {len(train_features)} 样本")
    print(f"- 验证数据: {len(val_features)} 样本")
    print(f"- 测试数据: {len(test_features)} 样本")
    
    return train_loader, val_loader, test_loader


def train_example_model(train_loader, val_loader, save_dir="./models"):
    """
    训练模型示例
    
    Args:
        train_loader: 训练数据加载器
        val_loader: 验证数据加载器
        save_dir: 模型保存目录
    
    Returns:
        model: 训练好的模型
        history: 训练历史
    """
    # 确保保存目录存在
    os.makedirs(save_dir, exist_ok=True)
    
    # 创建模型
    print("创建模型...")
    model = create_mahjong_model(
        input_dim=324,
        action_dim=6,
        hidden_dim=128,
        num_res_blocks=3
    )
    
    # 选择设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    model.to(device)
    
    # 训练模型
    print("开始训练...")
    start_time = datetime.now()
    
    history = train_model(
        model=model,
        train_dataset=train_loader.dataset,
        val_dataset=val_loader.dataset,
        num_epochs=5,  # 小批量测试，实际应用中应增加
        batch_size=32,
        lr=0.001,
        policy_weight=0.7,  # 策略损失权重
        value_weight=0.3,   # 价值损失权重
        device=device,
        log_dir="./logs",
        save_best=True,
        save_dir=save_dir
    )
    
    end_time = datetime.now()
    print(f"训练完成! 耗时: {end_time - start_time}")
    
    # 保存最终模型
    final_model_path = os.path.join(save_dir, "final_model.pt")
    torch.save(model.state_dict(), final_model_path)
    print(f"最终模型保存至: {final_model_path}")
    
    return model, history


def evaluate_example_model(model, test_loader):
    """
    评估模型性能示例
    
    Args:
        model: 训练好的模型
        test_loader: 测试数据加载器
    """
    print("评估模型性能...")
    
    # 选择设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    # 评估模型
    metrics = evaluate_model(
        model=model,
        dataset=test_loader.dataset,
        batch_size=32,
        device=device
    )
    
    # 打印评估结果
    print("模型评估结果:")
    print(f"- 总损失: {metrics['loss']:.4f}")
    print(f"- 策略损失: {metrics['policy_loss']:.4f}")
    print(f"- 价值损失: {metrics['value_loss']:.4f}")
    print(f"- 准确率: {metrics['accuracy']:.4f}")
    
    return metrics


def inference_example(model_path=None):
    """
    模型推理示例
    
    Args:
        model_path: 模型路径，如果为None则使用示例模型
    """
    print("模型推理示例...")
    
    # 选择设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 创建推理器
    if model_path and os.path.exists(model_path):
        print(f"从 {model_path} 加载模型")
        inference = MahjongInference(model_path=model_path, device=device)
    else:
        print("使用随机初始化的模型进行推理演示")
        # 创建一个随机初始化的模型用于演示
        model = create_mahjong_model(input_dim=324, action_dim=6)
        model.to(device)
        inference = MahjongInference(model=model, device=device)
    
    # 创建一个模拟的游戏状态
    game_state = MahjongGameState()
    
    # 创建特征工程器
    feature_engineer = FeatureEngineer()
    
    # 生成游戏状态的特征
    game_features = feature_engineer.create_game_features(game_state)
    
    # 进行推理
    action_probs, value = inference.predict(game_features)
    
    # 打印推理结果
    print("推理结果:")
    print(f"- 动作概率分布: {action_probs}")
    print(f"- 预测价值: {value:.4f}")
    print(f"- 推荐动作: {np.argmax(action_probs)}")
    
    # 使用决策器进行决策
    decision_maker = MahjongDecisionMaker(inference)
    best_action = decision_maker.make_decision(game_state)
    
    print(f"决策器选择的最佳动作: {best_action}")
    
    return inference, decision_maker


def export_model_example(model, export_dir="./exported_models"):
    """
    导出模型示例
    
    Args:
        model: 训练好的模型
        export_dir: 导出目录
    """
    print("导出模型...")
    
    # 确保导出目录存在
    os.makedirs(export_dir, exist_ok=True)
    
    # 导出为ONNX格式
    export_path = os.path.join(export_dir, "mahjong_model.onnx")
    export_trained_model(
        model=model,
        export_path=export_path,
        input_dim=324
    )
    
    print(f"模型导出至: {export_path}")
    
    # 导出为TorchScript格式
    ts_export_path = os.path.join(export_dir, "mahjong_model_script.pt")
    scripted_model = torch.jit.script(model)
    torch.jit.save(scripted_model, ts_export_path)
    
    print(f"TorchScript模型导出至: {ts_export_path}")
    
    return export_path


def main():
    """
    主函数，运行完整的示例流程
    """
    print("四川麻将AI模型使用示例")
    print("=" * 50)
    
    # 1. 准备数据
    train_loader, val_loader, test_loader = prepare_data()
    print()
    
    try:
        # 2. 训练模型
        model, history = train_example_model(train_loader, val_loader)
        print()
        
        # 3. 评估模型
        metrics = evaluate_example_model(model, test_loader)
        print()
        
        # 4. 导出模型
        export_path = export_model_example(model)
        print()
    except Exception as e:
        print(f"训练或评估过程中出错: {e}")
        print("使用随机初始化模型继续演示推理功能")
        model = None
        export_path = None
    
    # 5. 推理示例
    inference, decision_maker = inference_example(export_path)
    
    print()
    print("示例运行完成!")
    print("=" * 50)
    print("提示:")
    print("- 本示例使用随机生成的数据进行演示")
    print("- 实际应用中应使用真实的四川麻将游戏数据")
    print("- 为了更快的演示，训练轮数较少，实际应用中应增加训练轮数")
    print("- 可以调整模型参数、学习率等超参数以获得更好的性能")


if __name__ == "__main__":
    main()
